# pythonProject5/core/utils/audio.py
import logging
import wave
import contextlib
from pathlib import Path
from typing import Tuple, Optional, List, Generator
from datetime import datetime
from core.config import settings
import io
import tempfile
import os

# 引入必要的音频库
try:
    import librosa
except ImportError:
    librosa = None

# 尝试导入pydub用于音频转换
try:
    from pydub import AudioSegment
    PYDUB_AVAILABLE = True
except ImportError:
    PYDUB_AVAILABLE = False
    
logger = logging.getLogger(__name__)


def get_audio_info(audio_path: Path) -> Tuple[int, int, int, int]:
	"""
    获取音频文件信息
    :return: (帧数, 帧率, 通道数, 采样宽度)
    """
	if not audio_path.exists():
		raise FileNotFoundError(f"音频文件不存在: {audio_path}")

	try:
		with contextlib.closing(wave.open(str(audio_path), 'r')) as f:
			frames = f.getnframes()
			rate = f.getframerate()
			channels = f.getnchannels()
			width = f.getsampwidth()
			return frames, rate, channels, width
	except Exception as e:
		logger.error(f"读取音频文件失败 {audio_path}: {e}")
		raise


def get_duration(audio_path: Path) -> float:
	"""获取音频时长（秒）"""
	try:
		frames, rate, _, _ = get_audio_info(audio_path)
		return frames / rate
	except Exception:
		return 0.0


def validate_audio_file(audio_path: str | Path) -> tuple[bool, str]:
	"""
    验证音频文件是否有效
    """
	path = Path(audio_path)
	if not path.exists():
		return False, "文件不存在"
	if not path.is_file():
		return False, "路径不是文件"

	# 检查文件是否为音频文件
	try:
		# 使用 librosa 库验证音频文件
		if librosa:
			librosa.load(str(path), sr=None)
		else:
			logger.warning("librosa 库未安装，音频格式验证有限")
	except Exception as e:
		return False, f"无法解码音频: {e}"

	return True, "OK"


def create_upload_filename(user_id: int, prefix: str = "upload") -> str:
	"""生成上传文件名"""
	timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
	return f"{prefix}_{user_id}_{timestamp}.webm"


# ========================================================
# 音频流处理工具
# ========================================================

class AudioStreamProcessor:
    """音频流处理器，用于WebSocket实时音频处理"""
    
    def __init__(self, chunk_duration: float = 3.0, min_chunk_size: int = 1024):
        """
        初始化音频流处理器
        :param chunk_duration: 音频块时长（秒）
        :param min_chunk_size: 最小音频块大小（字节）
        """
        self.chunk_duration = chunk_duration
        self.min_chunk_size = min_chunk_size
        self.buffer = bytearray()
        self.last_process_time = datetime.now()
    
    def add_audio_data(self, audio_data: bytes) -> bool:
        """
        添加音频数据到缓冲区
        :param audio_data: 音频数据
        :return: 是否需要处理音频块
        """
        self.buffer.extend(audio_data)
        current_time = datetime.now()
        
        # 检查是否达到处理时间
        if (current_time - self.last_process_time).total_seconds() >= self.chunk_duration:
            if len(self.buffer) >= self.min_chunk_size:
                return True
        return False
    
    def get_audio_chunk(self) -> Optional[bytes]:
        """
        获取当前音频块并清空缓冲区
        :return: 音频数据，如果没有足够数据则返回None
        """
        if len(self.buffer) >= self.min_chunk_size:
            chunk = bytes(self.buffer)
            self.buffer.clear()
            self.last_process_time = datetime.now()
            return chunk
        return None
    
    def flush_buffer(self) -> Optional[bytes]:
        """
        强制处理缓冲区中的所有数据
        :return: 音频数据，如果没有数据则返回None
        """
        if len(self.buffer) > 0:
            chunk = bytes(self.buffer)
            self.buffer.clear()
            self.last_process_time = datetime.now()
            return chunk
        return None
    
    def get_buffer_size(self) -> int:
        """获取当前缓冲区大小"""
        return len(self.buffer)
    
    def reset(self):
        """重置处理器状态"""
        self.buffer.clear()
        self.last_process_time = datetime.now()


def validate_audio_chunk(audio_data: bytes, min_size: int = 512) -> tuple[bool, str]:
    """
    验证音频块是否有效
    :param audio_data: 音频数据
    :param min_size: 最小数据大小
    :return: (是否有效, 错误信息)
    """
    if len(audio_data) < min_size:
        return False, f"音频数据太小: {len(audio_data)} bytes"
    
    # 检查是否是WebM格式（简单检查）
    if len(audio_data) >= 4:
        # WebM文件头检查
        if audio_data[:4] == b'\x1a\x45\xdf\xa3':
            return True, "OK"
        # 其他格式可以在这里添加检查
    
    return True, "OK"


def create_temp_audio_file(audio_data: bytes, user_id: int, prefix: str = "chunk") -> str:
    """
    创建临时音频文件
    :param audio_data: 音频数据
    :param user_id: 用户ID
    :param prefix: 文件前缀
    :return: 临时文件路径
    """
    timestamp = int(datetime.utcnow().timestamp() * 1000)  # 毫秒级时间戳
    
    # 验证音频数据是否为有效的WebM格式
    is_valid_webm = False
    if len(audio_data) >= 4:
        # WebM文件以EBML标头开始 (\x1a\x45\xdf\xa3)
        if audio_data[:4] == b'\x1a\x45\xdf\xa3':
            is_valid_webm = True
    
    # 音频数据太小或不是有效的WebM格式
    if len(audio_data) < 1024:
        logger.warning(f"音频数据太小: {len(audio_data)}字节")
        # 如果数据太小，直接跳过写入
        if len(audio_data) < 512:
            logger.warning(f"音频数据太小，跳过写入: {len(audio_data)} 字节")
            return ""
    
    # 如果检测到无效的WebM格式，尝试转换
    if not is_valid_webm and len(audio_data) >= 1024:
        logger.info(f"检测到非WebM格式音频数据，尝试转换为WAV格式")
        try:
            # 创建一个临时的原始音频文件
            with tempfile.NamedTemporaryFile(delete=False, suffix='.raw') as raw_file:
                raw_path = raw_file.name
                raw_file.write(audio_data)
            
            # 使用pydub尝试转换音频（如果可用）
            if PYDUB_AVAILABLE:
                try:
                    # 尝试将原始数据作为PCM音频读取
                    # 使用多种采样率尝试
                    for sample_rate in [16000, 44100, 48000]:
                        try:
                            # 尝试作为PCM数据处理
                            audio = AudioSegment.from_raw(
                                raw_path, 
                                sample_width=2,  # 16位
                                frame_rate=sample_rate, 
                                channels=1  # 单声道
                            )
                            
                            # 创建临时WAV文件
                            filename = f"{prefix}_{user_id}_{timestamp}.wav"
                            temp_path = Path("temp_ws") / filename
                            
                            # 确保目录存在
                            temp_path.parent.mkdir(exist_ok=True)
                            
                            # 导出为WAV文件
                            audio.export(str(temp_path), format="wav")
                            logger.info(f"成功转换音频为WAV格式: {temp_path}, 采样率: {sample_rate}Hz")
                            
                            # 删除原始临时文件
                            os.unlink(raw_path)
                            
                            return str(temp_path)
                        except Exception as conv_e:
                            logger.debug(f"尝试采样率 {sample_rate}Hz 转换失败: {conv_e}")
                            continue
                    
                    # 如果所有采样率都失败
                    logger.warning("所有采样率尝试均失败，回退到原始WebM格式")
                except Exception as pydub_e:
                    logger.warning(f"pydub转换失败: {pydub_e}")
            else:
                logger.warning("pydub库不可用，无法转换音频格式")
            
            # 删除原始临时文件
            os.unlink(raw_path)
            
        except Exception as e:
            logger.error(f"音频转换过程中发生错误: {e}")
    
    # 使用原始WebM格式（如果转换失败或不需要转换）
    filename = f"{prefix}_{user_id}_{timestamp}.webm"
    temp_path = Path("temp_ws") / filename
    
    # 确保目录存在
    temp_path.parent.mkdir(exist_ok=True)
    
    # 写入音频数据
    with open(temp_path, "wb") as f:
        f.write(audio_data)
    
    # 检查写入的文件大小
    file_size = temp_path.stat().st_size
    if file_size < 512:
        logger.warning(f"写入的临时音频文件太小: {file_size} 字节")
    
    return str(temp_path)


def cleanup_temp_files(file_paths: List[str]):
    """
    清理临时文件
    :param file_paths: 文件路径列表
    """
    for file_path in file_paths:
        try:
            if Path(file_path).exists():
                Path(file_path).unlink()
        except Exception as e:
            logger.warning(f"无法删除临时文件 {file_path}: {e}")


def get_audio_format_info(audio_data: bytes) -> dict:
    """
    获取音频格式信息
    :param audio_data: 音频数据
    :return: 格式信息字典
    """
    info = {
        "size": len(audio_data),
        "format": "unknown",
        "valid": False
    }
    
    if len(audio_data) < 4:
        return info
    
    # WebM格式检查
    if audio_data[:4] == b'\x1a\x45\xdf\xa3':
        info["format"] = "webm"
        info["valid"] = True
    # 可以添加其他格式检查
    
    return info
